import polars as pl
import re
import os
import openpyxl
from openpyxl import load_workbook
from openpyxl.utils import get_column_letter
from openpyxl.styles import Font, PatternFill, Border, Side, Alignment, Color

def update_sheet_with_df(df, excel_file, sheet_name="Sheet1", start_cell="A1"):
    """
    将Polars DataFrame写入到现有Excel文件的指定sheet中，仅覆盖指定单元格范围的内容，其余内容保持不变
    如果sheet不存在，则自动创建
    
    参数:
    df: Polars DataFrame - 要写入的DataFrame
    excel_file: str - Excel文件路径
    start_cell: str - 起始单元格坐标，例如"B5"
    sheet_name: str - 工作表名称
    
    返回:
    int - 结束行的行号（1-based）
    """
    # 检查文件是否存在
    if not os.path.exists(excel_file):
        raise FileNotFoundError(f"文件 '{excel_file}' 不存在")
    
    # 解析起始单元格坐标
    match = re.match(r'([A-Z]+)(\d+)', start_cell)
    if not match:
        raise ValueError("起始单元格格式不正确，请使用类似'B5'的格式")
    
    start_col_str, start_row_str = match.groups()
    start_row = int(start_row_str)  # 保持1-based索引，因为openpyxl使用1-based
    
    # 将列字母转换为1-based索引
    start_col = 0
    for char in start_col_str:
        start_col = start_col * 26 + (ord(char) - ord('A') + 1)
    # 此时start_col已经是1-based
    
    # 加载现有工作簿
    wb = load_workbook(excel_file)
    
    # 如果sheet不存在，则创建它
    if sheet_name not in wb.sheetnames:
        ws = wb.create_sheet(sheet_name)
        print(f"工作表 '{sheet_name}' 不存在，已自动创建")
    else:
        ws = wb[sheet_name]
    
    # 创建格式对象
    from openpyxl.styles import Font, PatternFill, Border, Side, Alignment
    
    # 设置列名格式
    header_font = Font(bold=True)
    header_fill = PatternFill(start_color='D7E4BC', end_color='D7E4BC', fill_type='solid')
    
    # 设置边框
    border = Border(
        left=Side(style='thin'),
        right=Side(style='thin'),
        top=Side(style='thin'),
        bottom=Side(style='thin')
    )
    
    # 设置对齐方式
    alignment = Alignment(horizontal='center', vertical='center')
    
    # 写入列名
    for col_idx, col_name in enumerate(df.columns):
        col_letter = get_column_letter(start_col + col_idx)
        cell = ws[f"{col_letter}{start_row}"]
        cell.value = col_name
        
        # 应用列名样式
        cell.font = header_font
        cell.fill = header_fill
        cell.border = border
        cell.alignment = alignment
    
    # 写入数据
    for row_idx in range(len(df)):
        for col_idx, col_name in enumerate(df.columns):
            col_letter = get_column_letter(start_col + col_idx)
            cell = ws[f"{col_letter}{start_row + 1 + row_idx}"]
            
            # 直接获取值
            value = df[col_name][row_idx]
            cell.value = value
            
            # 应用数据单元格样式
            cell.border = border
    
    # 保存修改后的工作簿
    wb.save(excel_file)
    wb.close()
    
    # 计算并返回结束行号（1-based）
    end_row_1based = start_row + 1 + len(df)
    return end_row_1based

# 添加更新现有sheet的示例用法
def example_update_existing_sheet():
    """
    示例：更新现有Excel文件的指定sheet
    """
    # 创建要更新的DataFrame
    update_df = pl.DataFrame({
        "产品": ["笔记本电脑", "智能手机", "平板电脑"],
        "价格": [5999.99, 3999.50, 2999.00],
        "库存": [45, 120, 88]
    })
    
    # 假设我们要更新之前创建的example_output.xlsx文件
    output_file = "example_output.xlsx"
    sheet_name = "员工数据"  # 使用之前创建的sheet
    
    try:
        # 检查文件是否存在
        if not os.path.exists(output_file):
            print(f"文件 '{output_file}' 不存在，正在创建新文件和工作表...")
            # 创建新的工作簿和工作表
            wb = openpyxl.Workbook()
            if sheet_name != 'Sheet':
                ws = wb.create_sheet(sheet_name)
                # 删除默认的Sheet
                if 'Sheet' in wb.sheetnames:
                    wb.remove(wb['Sheet'])
            wb.save(output_file)
            wb.close()
            print(f"文件 '{output_file}' 和工作表 '{sheet_name}' 已创建")
        
        # 更新指定单元格范围
        end_row = update_sheet_with_df(update_df, output_file, sheet_name, "C12")
        print(f"\n数据已成功更新到Excel文件 '{output_file}'")
        print(f"工作表名称: {sheet_name}")
        print(f"更新起始位置: C12")
        print(f"更新结束行号: {end_row}")
        print(f"共更新 {len(update_df)} 行数据")
        print(f"共更新 {len(update_df.columns)} 列数据")
    except Exception as e:
        print(f"更新失败: {e}")

def write_text_with_style_to_cell(excel_file, sheet_name, cell_coord, text, style=None):
    """
    将指定文本和样式写入到Excel文件的指定单元格中
    
    参数:
    excel_file: str - Excel文件路径
    sheet_name: str - 工作表名称
    cell_coord: str - 单元格坐标，例如"B5"
    text: str - 要写入的文本内容
    style: dict - 可选的单元格样式字典，包含以下可选键:
        - font: Font对象
        - fill: PatternFill对象
        - border: Border对象
        - alignment: Alignment对象
        - number_format: str
    
    返回:
    bool - 操作是否成功
    """
    try:
        # 检查文件是否存在
        if not os.path.exists(excel_file):
            # 如果文件不存在，创建一个新的工作簿
            wb = openpyxl.Workbook()
            # 如果需要的sheet不是默认的，创建它
            if sheet_name != 'Sheet':
                ws = wb.create_sheet(sheet_name)
                # 删除默认的Sheet
                if 'Sheet' in wb.sheetnames:
                    wb.remove(wb['Sheet'])
            else:
                ws = wb.active
        else:
            # 加载现有工作簿
            wb = load_workbook(excel_file)
            # 确保sheet存在
            if sheet_name not in wb.sheetnames:
                ws = wb.create_sheet(sheet_name)
            else:
                ws = wb[sheet_name]
        
        # 写入文本到指定单元格
        ws[cell_coord] = text
        
        # 应用样式（如果提供）
        if style:
            cell = ws[cell_coord]
            
            if 'font' in style:
                cell.font = style['font']
            
            if 'fill' in style:
                cell.fill = style['fill']
            
            if 'border' in style:
                cell.border = style['border']
            
            if 'alignment' in style:
                cell.alignment = style['alignment']
            
            if 'number_format' in style:
                cell.number_format = style['number_format']
        
        # 保存并关闭工作簿
        wb.save(excel_file)
        wb.close()
        
        return True
    except Exception as e:
        print(f"写入单元格失败: {e}")
        return False

# 添加写入指定单元格的示例

def example_write_cell_with_style():
    """
    示例：将文本和样式写入到Excel文件的指定单元格
    """
    # 示例文件路径
    output_file = "example_output.xlsx"
    
    # 创建样式
    title_font = Font(bold=True, size=16, color="FF0000")
    title_fill = PatternFill(fill_type='solid')
    title_fill.bgColor = Color(rgb="FFD700")
    title_alignment = Alignment(horizontal='center', vertical='center')
    
    # 创建边框
    border = Border(
        left=Side(style='thin'),
        right=Side(style='thin'),
        top=Side(style='thin'),
        bottom=Side(style='thin')
    )
    
    # 定义样式字典
    title_style = {
        'font': title_font,
        'fill': title_fill,
        'border': border,
        'alignment': title_alignment
    }
    
    # 写入标题
    success = write_text_with_style_to_cell(
        output_file,
        "员工数据",
        "A1",
        "公司产品目录",
        title_style
    )
    
    if success:
        print(f"已成功将标题写入到 '{output_file}' 的 '员工数据' 工作表的 A1 单元格")
        
        # 再写入几个普通单元格
        write_text_with_style_to_cell(output_file, "员工数据", "A3", "产品ID")
        write_text_with_style_to_cell(output_file, "员工数据", "B3", "产品名称")
        write_text_with_style_to_cell(output_file, "员工数据", "C3", "价格")
        
        print("已成功写入表头信息")
    else:
        print("写入失败")

def write_list_to_excel(items_list, excel_file, sheet_name="Sheet1", start_row=1, start_col="A"):
    """
    将包含字典的列表内容写入到Excel文件中，根据每个字典的type字段决定写入方式
    
    参数:
    items_list: list - 包含字典的列表，每个字典格式为 {"type": string, "content": any}
                  type可以是 "string" 或 "dataframe"
    excel_file: str - Excel文件路径
    sheet_name: str - 工作表名称
    start_row: int - 起始行号（1-based）
    start_col: str - 起始列字母，例如 "A"
    
    返回:
    int - 最后的结束行号（1-based）
    """
    try:
        current_row = start_row
        
        # 遍历列表中的每个项目
        for idx, item in enumerate(items_list):
            if not isinstance(item, dict) or "type" not in item or "content" not in item:
                print(f"警告: 列表中第{idx+1}个元素格式不正确，跳过")
                continue
            
            item_type = item["type"]
            item_content = item["content"]
            
            # 构建当前单元格坐标
            current_cell = f"{start_col}{current_row}"
            
            if item_type == "string":
                # 写入字符串
                success = write_text_with_style_to_cell(
                    excel_file=excel_file,
                    sheet_name=sheet_name,
                    cell_coord=current_cell,
                    text=str(item_content)
                )
                if success:
                    print(f"已成功写入字符串到单元格 {current_cell}")
                    # 字符串写入后，行号加1
                    current_row += 1
                else:
                    print(f"写入字符串失败")
                    # 即使失败也尝试继续，行号加1
                    current_row += 1
            
            elif item_type == "dataframe":
                # 检查是否为Polars DataFrame
                if not isinstance(item_content, pl.DataFrame):
                    print(f"警告: 列表中第{idx+1}个元素的content不是Polars DataFrame，跳过")
                    continue
                
                # 写入DataFrame
                try:
                    end_row = update_sheet_with_df(
                        df=item_content,
                        excel_file=excel_file,
                        sheet_name=sheet_name,
                        start_cell=current_cell
                    )
                    print(f"已成功写入DataFrame，从 {current_cell} 开始，到第{end_row}行结束")
                    # 更新当前行号为DataFrame写入后的下一行
                    current_row = end_row + 1
                except Exception as e:
                    print(f"写入DataFrame失败: {e}")
                    # 即使失败也尝试继续，行号加1
                    current_row += 1
            
            else:
                print(f"警告: 列表中第{idx+1}个元素的type '{item_type}' 不支持，跳过")
        
        # 返回最终的结束行号
        return current_row
        
    except Exception as e:
        print(f"写入列表内容到Excel时出错: {e}")
        return -1

# 添加写入列表的示例
def example_write_list_to_excel():
    """
    示例：将包含字符串和DataFrame的列表写入到Excel文件
    """
    # 创建示例DataFrame
    df = pl.DataFrame({
        "产品": ["笔记本电脑", "智能手机", "平板电脑"],
        "价格": [5999.99, 3999.50, 2999.00],
        "库存": [45, 120, 88]
    })
    
    # 创建要写入的列表
    items_to_write = [
        {"type": "string", "content": "产品信息表"},
        {"type": "string", "content": ""},
        {"type": "string", "content": "以下是当前库存产品数据"},
        {"type": "dataframe", "content": df},
        {"type": "string", "content": "数据更新时间：2024-01-01"}
    ]
    
    # 示例文件路径
    output_file = "example_output.xlsx"
    sheet_name = "产品清单"
    
    # 删除现有文件以进行测试
    if os.path.exists(output_file):
        try:
            os.remove(output_file)
            print(f"已删除现有文件 '{output_file}'")
        except Exception as e:
            print(f"删除文件时出错: {e}")
            exit(-1)
    
    # 调用函数写入列表内容
    end_row = write_list_to_excel(
        items_list=items_to_write,
        excel_file=output_file,
        sheet_name=sheet_name,
        start_row=2,
        start_col="B"
    )
    
    if end_row != -1:
        print(f"\n列表内容已成功写入到Excel文件 '{output_file}'")
        print(f"工作表名称: {sheet_name}")
        print(f"起始位置: B2")
        print(f"结束行号: {end_row}")
    else:
        print("写入列表内容失败")

def safe_update_excel_file(df, excel_file, sheet_name="Sheet1", start_cell="A1"):
    """
    安全地将Polars DataFrame写入到Excel文件的指定sheet中，保留所有现有的sheet页
    如果文件不存在，则创建新文件
    如果sheet不存在，则创建新sheet
    不会删除任何现有的sheet页
    
    参数:
    df: Polars DataFrame - 要写入的DataFrame
    excel_file: str - Excel文件路径
    sheet_name: str - 工作表名称
    start_cell: str - 起始单元格坐标，例如"B5"
    
    返回:
    int - 结束行的行号（1-based）
    """
    try:
        # 检查文件是否存在
        if not os.path.exists(excel_file):
            print(f"文件 '{excel_file}' 不存在，正在创建新文件...")
            # 创建新的工作簿
            wb = openpyxl.Workbook()
            # 如果需要的sheet不是默认的，创建它（但保留默认Sheet）
            if sheet_name != 'Sheet':
                ws = wb.create_sheet(sheet_name)
            else:
                ws = wb.active
            wb.save(excel_file)
            wb.close()
            print(f"文件 '{excel_file}' 已创建")
        
        # 解析起始单元格坐标
        match = re.match(r'([A-Z]+)(\d+)', start_cell)
        if not match:
            raise ValueError("起始单元格格式不正确，请使用类似'B5'的格式")
        
        start_col_str, start_row_str = match.groups()
        start_row = int(start_row_str)  # 保持1-based索引
        
        # 将列字母转换为1-based索引
        start_col = 0
        for char in start_col_str:
            start_col = start_col * 26 + (ord(char) - ord('A') + 1)
        # 此时start_col已经是1-based
        
        # 加载工作簿
        wb = load_workbook(excel_file)
        
        # 如果sheet不存在，则创建它（保留其他所有sheet）
        if sheet_name not in wb.sheetnames:
            ws = wb.create_sheet(sheet_name)
            print(f"工作表 '{sheet_name}' 不存在，已自动创建")
        else:
            ws = wb[sheet_name]
        
        # 创建格式对象
        header_font = Font(bold=True)
        header_fill = PatternFill(start_color='D7E4BC', end_color='D7E4BC', fill_type='solid')
        
        border = Border(
            left=Side(style='thin'),
            right=Side(style='thin'),
            top=Side(style='thin'),
            bottom=Side(style='thin')
        )
        
        alignment = Alignment(horizontal='center', vertical='center')
        
        # 写入列名
        for col_idx, col_name in enumerate(df.columns):
            col_letter = get_column_letter(start_col + col_idx)
            cell = ws[f"{col_letter}{start_row}"]
            cell.value = col_name
            
            # 应用列名样式
            cell.font = header_font
            cell.fill = header_fill
            cell.border = border
            cell.alignment = alignment
        
        # 写入数据
        for row_idx in range(len(df)):
            for col_idx, col_name in enumerate(df.columns):
                col_letter = get_column_letter(start_col + col_idx)
                cell = ws[f"{col_letter}{start_row + 1 + row_idx}"]
                
                # 直接获取值
                value = df[col_name][row_idx]
                cell.value = value
                
                # 应用数据单元格样式
                cell.border = border
        
        # 保存修改后的工作簿
        wb.save(excel_file)
        wb.close()
        
        # 计算并返回结束行号（1-based）
        end_row_1based = start_row + 1 + len(df)
        return end_row_1based
    except Exception as e:
        print(f"安全更新Excel文件失败: {e}")
        return -1

# 修改现有的write_text_with_style_to_cell函数，移除删除默认sheet的逻辑
def write_text_with_style_to_cell(excel_file, sheet_name, cell_coord, text, style=None):
    """
    将指定文本和样式写入到Excel文件的指定单元格中
    保留所有现有的sheet页
    
    参数:
    excel_file: str - Excel文件路径
    sheet_name: str - 工作表名称
    cell_coord: str - 单元格坐标，例如"B5"
    text: str - 要写入的文本内容
    style: dict - 可选的单元格样式字典
    
    返回:
    bool - 操作是否成功
    """
    try:
        # 检查文件是否存在
        if not os.path.exists(excel_file):
            # 如果文件不存在，创建一个新的工作簿
            wb = openpyxl.Workbook()
            # 如果需要的sheet不是默认的，创建它（但保留默认Sheet）
            if sheet_name != 'Sheet':
                ws = wb.create_sheet(sheet_name)
            else:
                ws = wb.active
        else:
            # 加载现有工作簿
            wb = load_workbook(excel_file)
            # 确保sheet存在
            if sheet_name not in wb.sheetnames:
                ws = wb.create_sheet(sheet_name)
            else:
                ws = wb[sheet_name]
        
        # 写入文本到指定单元格
        ws[cell_coord] = text
        
        # 应用样式（如果提供）
        if style:
            cell = ws[cell_coord]
            
            if 'font' in style:
                cell.font = style['font']
            
            if 'fill' in style:
                cell.fill = style['fill']
            
            if 'border' in style:
                cell.border = style['border']
            
            if 'alignment' in style:
                cell.alignment = style['alignment']
            
            if 'number_format' in style:
                cell.number_format = style['number_format']
        
        # 保存并关闭工作簿
        wb.save(excel_file)
        wb.close()
        
        return True
    except Exception as e:
        print(f"写入单元格失败: {e}")
        return False

# 修改example_write_list_to_excel函数，移除删除现有文件的逻辑
def example_write_list_to_excel():
    """
    示例：将包含字符串和DataFrame的列表写入到Excel文件，保留所有现有sheet页
    """
    # 创建示例DataFrame
    df = pl.DataFrame({
        "产品": ["笔记本电脑", "智能手机", "平板电脑"],
        "价格": [5999.99, 3999.50, 2999.00],
        "库存": [45, 120, 88]
    })
    
    # 创建要写入的列表
    items_to_write = [
        {"type": "string", "content": "产品信息表"},
        {"type": "string", "content": ""},
        {"type": "string", "content": "以下是当前库存产品数据"},
        {"type": "dataframe", "content": df},
        {"type": "string", "content": "数据更新时间：2024-01-01"}
    ]
    
    # 示例文件路径
    output_file = "example_output.xlsx"
    sheet_name = "产品清单"
    
    # 不再删除现有文件，而是直接写入（保留所有现有sheet）
    print(f"将写入到文件 '{output_file}'，保留所有现有sheet页")
    
    # 调用函数写入列表内容
    end_row = write_list_to_excel(
        items_list=items_to_write,
        excel_file=output_file,
        sheet_name=sheet_name,
        start_row=2,
        start_col="B"
    )
    
    if end_row != -1:
        print(f"\n列表内容已成功写入到Excel文件 '{output_file}'")
        print(f"工作表名称: {sheet_name}")
        print(f"起始位置: B2")
        print(f"结束行号: {end_row}")
    else:
        print("写入列表内容失败")

# 添加一个示例函数来演示safe_update_excel_file的使用
def example_safe_update_excel():
    """
    示例：安全更新Excel文件，保留所有现有sheet页
    """
    # 创建要更新的DataFrame
    df1 = pl.DataFrame({
        "员工ID": [1001, 1002, 1003],
        "姓名": ["张三", "李四", "王五"],
        "部门": ["技术部", "市场部", "财务部"]
    })
    
    df2 = pl.DataFrame({
        "订单号": ["ORD-001", "ORD-002", "ORD-003"],
        "金额": [1000.00, 2500.50, 800.75],
        "状态": ["已完成", "处理中", "已取消"]
    })
    
    # 示例文件路径
    output_file = "example_output.xlsx"
    
    # 写入第一个sheet
    print("正在写入'员工信息'工作表...")
    end_row1 = safe_update_excel_file(df1, output_file, "员工信息", "A1")
    
    # 写入第二个sheet
    print("\n正在写入'订单信息'工作表...")
    end_row2 = safe_update_excel_file(df2, output_file, "订单信息", "A1")
    
    # 再次更新第一个sheet（添加一些额外数据）
    df3 = pl.DataFrame({
        "员工ID": [1004, 1005],
        "姓名": ["赵六", "钱七"],
        "部门": ["人力资源", "销售部"]
    })
    print("\n正在更新'员工信息'工作表...")
    end_row3 = safe_update_excel_file(df3, output_file, "员工信息", f"A{end_row1 + 2}")
    
    if end_row1 != -1 and end_row2 != -1 and end_row3 != -1:
        print(f"\n所有数据已成功写入到Excel文件 '{output_file}'")
        print(f"'员工信息'工作表结束行: {end_row3}")
        print(f"'订单信息'工作表结束行: {end_row2}")
    else:
        print("写入数据失败")

if __name__ == "__main__":
    # 运行更新现有sheet示例
    # example_update_existing_sheet()  # 运行更新现有sheet示例
    
    # 运行写入指定单元格的示例
    # example_write_cell_with_style()  # 运行写入单元格示例
    
    # 运行写入列表的示例
    # example_write_list_to_excel()  # 运行写入列表示例
    
    # 运行安全更新Excel的示例（保留所有sheet）
    example_safe_update_excel()  # 运行安全更新示例